using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using UnityEditor;
using UnityEditor.Rendering;
using UnityEngine;
using Debug = UnityEngine.Debug;

namespace Mars.ShaderAnalysis
{
    public class CompileVariantAnalysis
    {
        public Shader shader;
        public int subShaderIndex;
        public int passIndex;
        public string[] keywords;
        public ShaderType shaderType;
        public ShaderCompilerPlatform shaderCompilerPlatform;
        public BuildTarget buildTarget;

        //分析数据
        public uint totalCount = 0;
        public uint forCount = 0;
        public uint ifCount = 0;
        public uint ternaryCount = 0;
        public uint dotCount = 0;
        public uint minMaxCount = 0;
        public uint sinCosCount = 0;
        public uint specialOperationCount = 0;
        public uint discardCount = 0;
        public uint clampCount = 0;
        public uint absCount = 0;
        public uint textureCount = 0;

        public uint floatCount = 0;
        public uint halfCount = 0;
        public uint intCount = 0;
        public uint uintCount = 0;
        public uint boolCount = 0;
        public uint unknownCount = 0;

        private void ResetAnalysis()
        {
            totalCount = 0;
            forCount = 0;
            ifCount = 0;
            ternaryCount = 0;
            dotCount = 0;
            minMaxCount = 0;
            sinCosCount = 0;
            specialOperationCount = 0;
            discardCount = 0;
            clampCount = 0;
            absCount = 0;
            textureCount = 0;

            floatCount = 0;
            halfCount = 0;
            intCount = 0;
            uintCount = 0;
            boolCount = 0;
            unknownCount = 0;
        }

        public void SetData(Shader shader, int subShaderIndex, int passIndex, string[] keywords, ShaderType shaderType, ShaderCompilerPlatform shaderCompilerPlatform, BuildTarget buildTarget)
        {
            this.shader = shader;
            this.subShaderIndex = subShaderIndex;
            this.passIndex = passIndex;
            this.keywords = keywords;
            this.shaderType = shaderType;
            this.shaderCompilerPlatform = shaderCompilerPlatform;
            this.buildTarget = buildTarget;
        }

        /// <summary>
        /// 执行分析
        /// </summary>
        public string Execute()
        {
            string filePath = Compile(shader, subShaderIndex, passIndex, keywords, shaderType, shaderCompilerPlatform, buildTarget);
            if (string.IsNullOrEmpty(filePath))
            {
                return filePath;
            }

            ResetAnalysis();
            int firstOutLine = -1;
            int firstMainLine = -1;

            HashSet<string> halfHash = new HashSet<string>();
            HashSet<string> floatHash = new HashSet<string>();
            HashSet<string> intHash = new HashSet<string>();
            HashSet<string> uintHash = new HashSet<string>();
            HashSet<string> boolHash = new HashSet<string>();
            string[] allLines = File.ReadAllLines(filePath);
            for (int i = 0; i < allLines.Length; i++)
            {
                string line = allLines[i];

                if (line == "void main()")
                {
                    firstMainLine = i;
                }

                if (line.Contains("in "))
                {
                    firstOutLine = i + 1;
                    continue;
                }

                if (line.Contains("out "))
                {
                    firstOutLine = i;
                }

                //属性类型：输出属性与主函数之间
                if (firstOutLine >= 0 && firstMainLine < 0)
                {
                    //临时变量
                    if (line.Contains("u_xlat"))
                    {
                        string valName = line.Split("u_xlat")[1].Replace(";", "");
                        valName = "u_xlat" + valName;
                        if (line.Contains("u_xlatu"))
                        {
                            uintHash.Add(valName);
                        }
                        else if (line.Contains("u_xlati"))
                        {
                            intHash.Add(valName);
                        }
                        else if (line.Contains("u_xlatb"))
                        {
                            boolHash.Add(valName);
                        }
                        else if (line.Contains("mediump"))
                        {
                            halfHash.Add(valName);
                        }
                        else
                        {
                            floatHash.Add(valName);
                        }
                    }
                    //输出变量
                    else if (line.Contains("out "))
                    {
                        if (line.Contains(" vs_"))
                        {
                            //顶点输出
                            string valName = line.Split(" vs_")[1].Replace(";", "");
                            valName = "vs_" + valName;
                            CheckOutputVal(line, valName, uintHash, intHash, boolHash, halfHash, floatHash);
                        }
                        else if (line.Contains(" SV_"))
                        {
                            //片元输出
                            string valName = line.Split(" SV_")[1].Replace(";", "");
                            valName = "SV_" + valName;
                            CheckOutputVal(line, valName, uintHash, intHash, boolHash, halfHash, floatHash);
                        }
                        else
                        {
                            Debug.LogError("未知输出变量 : " + line);
                        }
                    }
                    else
                    {
                        Debug.LogError("未知变量 : " + line);
                    }
                }

                //ALU计算 :找到了主函数
                if (firstMainLine >= 0 && i > firstMainLine)
                {
                    //去掉字符串里的空格
                    line = line.Trim().Replace(" ", "");
                    if (IsALU(line))
                    {
                        totalCount++;
                        //指令类型判断
                        if (line.Contains("for("))
                        {
                            forCount++;
                        }

                        if (line.Contains("if("))
                        {
                            ifCount++;
                        }

                        if (line.Contains("?") && line.Contains(":"))
                        {
                            ternaryCount++;
                        }

                        if (line.Contains("dot("))
                        {
                            dotCount++;
                        }

                        if (line.Contains("clamp("))
                        {
                            clampCount++;
                        }

                        if (line.Contains("abs("))
                        {
                            absCount++;
                        }

                        if (line.Contains("min(") || line.Contains("max("))
                        {
                            minMaxCount++;
                        }

                        if (line.Contains("sin(") || line.Contains("cos("))
                        {
                            sinCosCount++;
                        }

                        if (line.Contains("texture(") || line.Contains("textureLod("))
                        {
                            textureCount++;
                        }

                        if (line.Contains("sqrt(") || line.Contains("inversesqrt(") || line.Contains("fract(") || line.Contains("log2(") || line.Contains("exp2("))
                        {
                            specialOperationCount++;
                        }

                        if (line.Contains("discard;"))
                        {
                            discardCount++;
                        }

                        //指令精度判断
                        string valName = line.Split('=')[0].Split('.')[0];
                        if (floatHash.Contains(valName))
                        {
                            floatCount++;
                        }
                        else if (halfHash.Contains(valName))
                        {
                            halfCount++;
                        }
                        else if (intHash.Contains(valName))
                        {
                            intCount++;
                        }
                        else if (uintHash.Contains(valName))
                        {
                            uintCount++;
                        }
                        else if (boolHash.Contains(valName))
                        {
                            boolCount++;
                        }
                        else
                        {
                            //直接枚举一些例外的存在吧
                            if (line.StartsWith("for(uint") || line.StartsWith("uvec") || line.StartsWith("uint"))
                            {
                                uintCount++;
                            }
                            else if (line.StartsWith("for(int") || line.StartsWith("ivec") || line.StartsWith("int"))
                            {
                                intCount++;
                            }
                            else if (line.StartsWith("if(") || line.StartsWith("bvec") || line.StartsWith("bool"))
                            {
                                boolCount++;
                            }
                            else if (line.StartsWith("gl_Position") || line.StartsWith("vec") || line.StartsWith("float") ||
                                     line.StartsWith("ImmCB_") || line.StartsWith("hlslcc_movcTemp") || line.StartsWith("gl_FragDepth") ||
                                     line.StartsWith("u_xlat0"))
                            {
                                floatCount++;
                            }
                            else if (line.StartsWith("mediump"))
                            {
                                halfCount++;
                            }
                            else
                            {
                                unknownCount++;
                                Debug.LogError("未知指令 : " + line);
                            }
                        }
                    }
                    else
                    {
                        //Debug.LogError(line);
                    }
                }
            }

            return filePath;
        }

        /// <summary>
        /// 检查输出变量属于哪种精度类型
        /// </summary>
        private void CheckOutputVal(string line, string valName, HashSet<string> uintHash, HashSet<string> intHash, HashSet<string> boolHash, HashSet<string> halfHash, HashSet<string> floatHash)
        {
            if (line.Contains("uvec") || line.Contains("uint"))
            {
                uintHash.Add(valName);
            }
            else if (line.Contains("ivec") || line.Contains("int"))
            {
                intHash.Add(valName);
            }
            else if (line.Contains("bvec") || line.Contains("bool"))
            {
                boolHash.Add(valName);
            }
            else if (line.Contains("vec") || line.Contains("float"))
            {
                if (line.Contains("mediump"))
                {
                    halfHash.Add(valName);
                }
                else
                {
                    floatHash.Add(valName);
                }
            }
            else
            {
                Debug.LogError("未知输出变量 : " + line);
            }
        }

        /// <summary>
        /// 是否为指令字符串
        /// </summary>
        private bool IsALU(string line)
        {
            return !(line == "{" || line == "}" || line == "}else{" || line == "return;" || line == ";") &&
                   !string.IsNullOrEmpty(line);
        }

        /// <summary>
        /// 编译变体
        /// </summary>
        public static string Compile(Shader shader, int subShaderIndex, int passIndex, string[] keywords,
            ShaderType shaderType, ShaderCompilerPlatform shaderCompilerPlatform, BuildTarget buildTarget)
        {
            ShaderData shaderData = ShaderUtil.GetShaderData(shader);
            ShaderData.Subshader subshader = shaderData.GetSubshader(subShaderIndex);
            ShaderData.Pass pass = subshader.GetPass(passIndex);
            ShaderData.VariantCompileInfo variantCompileInfo = pass.CompileVariant(shaderType, keywords, shaderCompilerPlatform, buildTarget, true);
            string folderPath = EditorTools.projectPath + "CompileShaders";
            if (!Directory.Exists(folderPath))
            {
                Directory.CreateDirectory(folderPath);
            }

            folderPath += "/" + buildTarget;
            if (!Directory.Exists(folderPath))
            {
                Directory.CreateDirectory(folderPath);
            }

            folderPath += "/" + shaderCompilerPlatform;
            if (!Directory.Exists(folderPath))
            {
                Directory.CreateDirectory(folderPath);
            }

            folderPath += "/" + shaderType;
            if (!Directory.Exists(folderPath))
            {
                Directory.CreateDirectory(folderPath);
            }

            folderPath += "/" + shader.name.Replace("/", "_");
            if (!Directory.Exists(folderPath))
            {
                Directory.CreateDirectory(folderPath);
            }

            folderPath += "/" + "SubShader_" + subShaderIndex;
            if (!Directory.Exists(folderPath))
            {
                Directory.CreateDirectory(folderPath);
            }

            folderPath += "/" + "Pass_" + passIndex;
            if (!Directory.Exists(folderPath))
            {
                Directory.CreateDirectory(folderPath);
            }

            //显示拼接字符串
            string showKeywordString = "";
            for (int k = 0; k < keywords.Length; k++)
            {
                showKeywordString += keywords[k];
                if (k != keywords.Length - 1)
                {
                    showKeywordString += " ";
                }
            }

            string filePath = folderPath + "/" + showKeywordString.GetHashCode() + ".hlsl";
            File.WriteAllBytes(filePath, variantCompileInfo.ShaderData);

            return filePath;
        }

        /// <summary>
        /// 显示GUI
        /// </summary>
        public void OnGUI()
        {
            using (new EditorGUILayout.VerticalScope("FrameBox"))
            {
                using (new EditorGUILayout.HorizontalScope())
                {
                    EditorGUILayout.LabelField((shaderType == ShaderType.Vertex ? "顶点" : "片元") + "着色器", EditorStyles.boldLabel);
                    using (new EditorGUI.DisabledGroupScope(shader == null))
                    {
                        using (new EditorGUILayoutUtility.GUIColorScope(Color.magenta))
                        {
                            if (GUILayout.Button("使用malioc分析", EditorStyles.miniButton))
                            {
                                //编译Shader
                                string filePath = Execute();
                                //打开文本文件
                                string fileContent = File.ReadAllText(filePath);
                                //工具要 310 es 才能分析成功
                                fileContent = fileContent.Replace("#version 300 es", "#version 310 es");
                                filePath = filePath.Replace(".hlsl", shaderType == ShaderType.Vertex ? ".vert" : ".frag");
                                File.WriteAllText(filePath, fileContent);
                                //Debug.LogError(filePath);
                                //TCEditorUtility.Cmd("malioc \"" + filePath + "\"");
                                string args = "malioc \"" + filePath + "\"";
                                ProcessStartInfo pStartInfo = new ProcessStartInfo("cmd.exe");
                                pStartInfo.Arguments = "/k " + args;
                                Debug.Log(args);
                                Process.Start(pStartInfo);
                            }
                        }

                        using (new EditorGUILayoutUtility.GUIColorScope(Color.cyan))
                        {
                            if (GUILayout.Button("查看", EditorStyles.miniButton, GUILayout.Width(50f)))
                            {
                                //编译Shader
                                string filePath = Execute();
                                //打开文本文件
                                EditorUtility.OpenWithDefaultApp(filePath);
                            }
                        }
                    }
                }

                using (new EditorGUILayout.HorizontalScope())
                {
                    EditorGUILayout.LabelField(EditorGUIUtility.TrTextContent("编译平台 : "), GUILayout.Width(60f));
                    EditorGUILayout.LabelField(EditorGUIUtility.TrTextContent(shaderCompilerPlatform.ToString()));
                }

                using (new EditorGUILayout.HorizontalScope())
                {
                    EditorGUILayout.LabelField(EditorGUIUtility.TrTextContent("总指令数 : ", "总指令数建议不超过400"), GUILayout.Width(60f));
                    using (new EditorGUILayoutUtility.GUIColorScope(totalCount >= 400 ? Color.red : Color.green))
                    {
                        EditorGUILayout.LabelField(EditorGUIUtility.TrTextContent(totalCount.ToString()));
                    }
                }

                GUILayout.Space(2f);

                using (new EditorGUILayoutUtility.GUIColorScope(Color.white * 0.9f))
                {
                    using (new EditorGUILayout.VerticalScope("FrameBox"))
                    {
                        //float超过一定限制，报个红色
                        bool isRed = totalCount >= 100 && (float)floatCount / totalCount > 0.40f;
                        using (new EditorGUILayoutUtility.GUIColorScope(isRed ? Color.red : GUI.color))
                        {
                            EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("Float指令数", "总指令数100以上时，float指令比例建议不超过40%"),
                                floatCount + " (" + ((float)floatCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        }

                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("Half指令数"),
                            halfCount + " (" + ((float)halfCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("int指令数"),
                            intCount + " (" + ((float)intCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("uint指令数"),
                            uintCount + " (" + ((float)uintCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("bool指令数"),
                            boolCount + " (" + ((float)boolCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        if (unknownCount > 0)
                        {
                            EditorGUILayoutUtility.DrawLine();
                            using (new EditorGUILayoutUtility.GUIColorScope(Color.red))
                            {
                                EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("未知指令数", "请查看Log"),
                                    unknownCount + " (" + ((float)unknownCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                            }
                        }
                    }

                    GUILayout.Space(3f);
                    using (new EditorGUILayout.VerticalScope("FrameBox"))
                    {
                        //采样次数超过一定限制，报个红色
                        bool isRed = textureCount > 9;
                        using (new EditorGUILayoutUtility.GUIColorScope(isRed ? Color.red : GUI.color))
                        {
                            EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("texture指令数", "贴图采样建议不超过9次"),
                                textureCount + " (" + ((float)textureCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        }

                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("特殊运算指令数", "包含 sqrt(), inversesqrt(), fract(), log2(), exp2()"),
                            specialOperationCount + " (" + ((float)specialOperationCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("dot指令数"),
                            dotCount + " (" + ((float)dotCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("clamp指令数"),
                            clampCount + " (" + ((float)clampCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("minMax指令数"),
                            minMaxCount + " (" + ((float)minMaxCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("abs指令数"),
                            absCount + " (" + ((float)absCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("sinCos指令数"),
                            sinCosCount + " (" + ((float)sinCosCount / totalCount * 100).ToString("0.00") + "%)", EditorStyles.label);
                    }

                    GUILayout.Space(3f);
                    using (new EditorGUILayout.VerticalScope("FrameBox"))
                    {
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("for循环个数"), forCount.ToString(), EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("if个数"), ifCount.ToString(), EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("三目运算个数"), ternaryCount.ToString(), EditorStyles.label);
                        EditorGUILayoutUtility.DrawLine();
                        EditorGUILayout.TextField(EditorGUIUtility.TrTextContent("discard个数"), discardCount.ToString(), EditorStyles.label);
                    }
                }
            }

            EditorGUILayout.Space();
        }
    }
}